import os,json
import numpy as np
import cv2
from PIL import Image
class Sketlonize():
    def __init__(self,vascular_threshold) :
        self.vascular_threshold=vascular_threshold
    def optic_disc_vascular(self, vascular_path, optic_disc_coor, radius=5):
        '''Sets the area around optic_disc_coor with the specified radius to 0 to clear that region'''
        # Read the image in grayscale
        img = cv2.imread(vascular_path, cv2.IMREAD_GRAYSCALE)

        if img is None:
            raise FileNotFoundError(f"Image at path {vascular_path} could not be loaded.")

        # Create a mask with the same dimensions as the image, initialized to 1 (all values)
        mask = np.ones_like(img, dtype=np.uint8)

        # Set the circular region around the optic_disc_coor to 0
        cv2.circle(mask, optic_disc_coor, radius, 0, -1)

        # Apply the mask to the image to clear the region
        img = img * mask

        return img
    def sketlonize(self,vascular_path,optic_disc_coor,min_size=20,raduis=5):
        img=self.optic_disc_vascular(vascular_path,optic_disc_coor,raduis)
        img=img/255
        img = np.where(img > self.vascular_threshold, 255, 0).astype(np.uint8) # 阈值化后将像素值设为 255（前景）和 0（背景）
        # 骨架化
        img=img.astype(np.uint8)
        skeleton = cv2.ximgproc.thinning(img)
        return self.filer_sketlon(skeleton,min_size)
    def filer_sketlon(self,skeleton,min_size=20):
        '''
        使用 OpenCV 过滤掉骨架图中小于指定最小大小的连通区域。

        输入:
            skeleton (ndarray): 输入的二值骨架图，0 表示背景，1 或 255 表示前景。
            min_size (int): 最小保留的连通区域大小（以像素为单位）。

        输出:
            filtered_skeleton (ndarray): 过滤后的骨架图，只有符合条件的区域被保留。
        '''
        # 转换为适合连通区域分析的格式
        skeleton = skeleton.astype(np.uint8)

        # 找到连通区域
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(skeleton, connectivity=8)

        # 初始化一个空的骨架图
        filtered_skeleton = np.zeros_like(skeleton)

        # 遍历每个区域，保留符合条件的区域
        vascular_cnt=1
        vascular_label_img=np.zeros_like(skeleton)
        for label in range(1, num_labels):  # 跳过背景标签（0）
            area = stats[label, cv2.CC_STAT_AREA]  # 获取区域的面积
            if area >= min_size:
                filtered_skeleton[labels == label] = 255  # 保留区域并设置为 255（前景）
                vascular_label_img[labels == label] = vascular_cnt
                vascular_cnt+=1
                
        return filtered_skeleton,vascular_label_img,vascular_cnt-1
    def visual_sketlon_above(self, skeleton_image, src_image_path, save_path, thing=1):
        '''
        在原图上叠加骨架图，骨架用黄色显示，并加粗到thing像素，原图可能是单通道图像，如果是，需要判断并转化。

        '''
        # 读取原图
        src_image = cv2.imread(src_image_path)
        
        # 检查图像是否为单通道，如果是，将其转换为三通道
        if len(src_image.shape) == 2 or src_image.shape[2] == 1:
            src_image = cv2.cvtColor(src_image, cv2.COLOR_GRAY2BGR)
        
        # 创建一个掩膜，将骨架叠加到原图上
        mask = np.zeros_like(src_image)
        mask[skeleton_image == 255] = [0, 255, 255]  # 用黄色（BGR格式）显示骨架
        
        # 使用膨胀操作加粗骨架线条
        kernel = np.ones((thing, thing), np.uint8)
        mask = cv2.dilate(mask, kernel, iterations=1)
        
        # 将掩膜叠加到原图上
        overlay = cv2.addWeighted(src_image, 1.0, mask, 0.5, 0)
        
        # 保存结果
        cv2.imwrite(save_path, overlay)
    def visual_raw_skeleton(self, skeleton, save_path, thing=3):
        '''
        直接可视化骨架图，并加粗到thing像素。
        输入的骨架图是单通道的二值图像（0和255）。
        '''
        # 创建一个三通道图像用于显示骨架
        skeleton_color = np.zeros((skeleton.shape[0], skeleton.shape[1], 3), dtype=np.uint8)
        skeleton_color[skeleton == 255] = [0, 255, 255]  # 用黄色显示骨架
        if thing>1:
            # 使用膨胀操作加粗骨架线条
            kernel = np.ones((thing, thing), np.uint8)
            skeleton_color = cv2.dilate(skeleton_color, kernel, iterations=1)

        # 保存结果
        cv2.imwrite(save_path, skeleton_color)

if __name__ == '__main__':
    pass